import pandas as pd
import geopandas as gpd
import folium
from folium.plugins import TimestampedGeoJson
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import contextily as ctx
import os
from PIL import Image
import imageio
from datetime import datetime
import numpy as np


# 1. 数据准备（假设CSV文件包含timestamp, lat, lon列）
def load_gps_data(file_path):
    """加载GPS数据并预处理"""
    df = pd.read_csv(file_path)

    # 转换时间格式（如果时间戳是字符串）
    df['timestamp'] = pd.to_datetime(df['timestamp'])

    # 排序数据
    df = df.sort_values('timestamp').reset_index(drop=True)

    # 计算点之间的时间间隔（用于动画速度）
    df['time_diff'] = df['timestamp'].diff().dt.total_seconds().fillna(0)

    return df


# 2. 创建地图动画帧
def create_animation_frames(gps_data, output_dir='frames'):
    """生成动画帧图像"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 创建GeoDataFrame
    geometry = gpd.points_from_xy(gps_data.lon, gps_data.lat)
    gdf = gpd.GeoDataFrame(gps_data, geometry=geometry, crs="EPSG:4326")

    # 转换为Web墨卡托投影（用于底图）
    gdf = gdf.to_crs(epsg=3857)

    # 计算地图边界
    minx, miny, maxx, maxy = gdf.total_bounds
    margin = 0.1 * max((maxx - minx), (maxy - miny))

    # 创建基础地图
    fig, ax = plt.subplots(figsize=(10, 8))

    # 初始化轨迹和当前位置点
    line, = ax.plot([], [], 'b-', linewidth=2, alpha=0.7)
    point, = ax.plot([], [], 'ro', markersize=8)

    def init():
        ax.set_xlim(minx - margin, maxx + margin)
        ax.set_ylim(miny - margin, maxy + margin)
        ax.set_axis_off()
        return line, point

    def update(frame):
        # 更新轨迹
        x = gdf.geometry.x.iloc[:frame + 1]
        y = gdf.geometry.y.iloc[:frame + 1]
        line.set_data(x, y)

        # 更新当前位置
        if frame < len(gdf):
            point.set_data([gdf.geometry.x.iloc[frame]], [gdf.geometry.y.iloc[frame]])

        # 添加底图
        ctx.add_basemap(ax, source=ctx.providers.OpenStreetMap.Mapnik)

        # 添加标题（显示时间）
        timestamp = gps_data.timestamp.iloc[frame]
        ax.set_title(f"GPS Track - {timestamp.strftime('%Y-%m-%d %H:%M:%S')}", fontsize=14)

        # 保存当前帧
        plt.savefig(f"{output_dir}/frame_{frame:04d}.png", dpi=120, bbox_inches='tight')
        return line, point

    # 创建动画（但不实际显示）
    anim = FuncAnimation(fig, update, frames=len(gdf), init_func=init, blit=True)

    # 保存最后一帧作为完整轨迹
    plt.close()
    return len(gdf)


# 3. 从帧创建视频
def create_video_from_frames(frame_count, output_file='gps_track.mp4', fps=10):
    """将帧图像合成为视频"""
    # 确定帧率（根据数据点时间间隔调整）
    with imageio.get_writer(output_file, fps=fps) as writer:
        for i in range(frame_count):
            image = imageio.imread(f"frames/frame_{i:04d}.png")
            writer.append_data(image)

    print(f"视频已保存至: {output_file}")


# 4. 可选：创建交互式HTML地图
def create_interactive_map(gps_data, output_file='gps_map.html'):
    """创建带时间轴的交互式地图"""
    # 创建基础地图
    m = folium.Map(
        location=[gps_data.lat.mean(), gps_data.lon.mean()],
        zoom_start=13,
        tiles='OpenStreetMap'
    )

    # 准备GeoJSON数据
    features = []
    for idx, row in gps_data.iterrows():
        feature = {
            'type': 'Feature',
            'geometry': {
                'type': 'Point',
                'coordinates': [row.lon, row.lat]
            },
            'properties': {
                'time': row.timestamp.isoformat(),
                'popup': f"时间: {row.timestamp}<br>经度: {row.lon:.6f}<br>纬度: {row.lat:.6f}",
                'icon': 'circle',
                'iconstyle': {
                    'fillColor': '#ff0000',
                    'fillOpacity': 0.8,
                    'stroke': False,
                    'radius': 5
                }
            }
        }
        features.append(feature)

    # 添加时间轴轨迹
    TimestampedGeoJson(
        {'type': 'FeatureCollection', 'features': features},
        period='PT1M',  # 时间间隔显示
        add_last_point=True,
        auto_play=True,
        loop=False,
        max_speed=10,
        loop_button=True,
        date_options='YYYY-MM-DD HH:mm:ss',
        time_slider_drag_update=True
    ).add_to(m)

    # 保存HTML地图
    m.save(output_file)
    print(f"交互式地图已保存至: {output_file}")


# 主函数
def main():
    # 配置参数
    GPS_FILE = './data/gps_data.csv'  # 替换为你的GPS数据文件
    OUTPUT_VIDEO = 'gps_track_animation.mp4'
    OUTPUT_HTML = 'interactive_gps_map.html'
    FPS = 15  # 视频帧率

    # 1. 加载数据
    print("加载GPS数据...")
    gps_data = load_gps_data(GPS_FILE)
    print(f"已加载 {len(gps_data)} 个GPS数据点")
    print(f"时间范围: {gps_data.timestamp.min()} 到 {gps_data.timestamp.max()}")

    # 2. 创建动画帧
    print("生成动画帧...")
    frame_count = create_animation_frames(gps_data)
    print(f"已生成 {frame_count} 个动画帧")

    # # 3. 创建视频
    # print("创建视频...")
    # create_video_from_frames(frame_count, OUTPUT_VIDEO, FPS)

    # 4. 创建交互式地图
    print("创建交互式地图...")
    create_interactive_map(gps_data, OUTPUT_HTML)

    print("处理完成!")


if __name__ == "__main__":
    main(i love you)
